# coding: utf-8
import math

# from quant_researcher.quant.datasource_fetch.index_api.index_price_related import C_DATE_ORIGIN, get_index_monthly_return
from quant_researcher.quant.datasource_fetch.index_api.index_price_related import get_index_return
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool import hammer

import pandas as pd

C_DATE_ORIGIN = 'trade_date'


def main(index_names, risk_free, begin_date, end_date):
    LOG.info(f'Working from {begin_date} to {end_date} for {index_names}')

    m_risk_free = math.pow(1+risk_free, 1/12.0) - 1
    LOG.info(f'Risk free rate: {risk_free}，分解到月：{m_risk_free}')

    # store information about each index selected as market standard
    market_represents = {}

    def get_market_info(code):
        if not market_represents.__contains__(code):
            # 提醒：get_index_monthly_return取的时间区间为[begin_date, end_date],
            #      get_index_return取得时间区间为[begin_date[0:7], end_date[0:7]]
            # data_df = get_index_monthly_return(begin_date, end_date, code)
            data_df = get_index_return(
                index_code=code,
                start_date=begin_date,
                end_date=end_date,
            )
            data_df['end_date'] = data_df['end_date'].apply(lambda x: x[0:7])
            data_df = data_df.rename(columns={'end_date': 'tj', 'daily_return': 'i_price'})
            if data_df.shape[0] == 0:
                raise ValueError(f'没找到 {code} 从 {begin_date} 到 {end_date} 的行情')

            mean = data_df[code].mean()
            variance = data_df[code].var()
            LOG.info('%s的月频数据, mean: %s, variance: %s', code, mean, variance)
            market_represents[code] = {
                'data': data_df, 'mean': mean, 'variance': variance
            }
        return market_represents[code]

    def get_beta(m_code, index_data):
        market_info = get_market_info(m_code)
        market_data = market_info['data']
        market_variance = market_info['variance']

        data = pd.merge(market_data, index_data, how='outer', on=C_DATE_ORIGIN)
        del data[C_DATE_ORIGIN]
        hammer.slim(data, a_f=1)

        cov = data.cov().values[0][1]
        beta_here = cov / market_variance
        LOG.info('Code: %s, cov: %s, market_variance: %s, beta: %s',
                 m_code, cov, market_variance, beta_here)
        return beta_here, market_info['mean']

    ans_dict = {}  # final results
    for index_info in index_names:
        index_code = index_info[0]
        index_market_standard = index_info[1]
        LOG.info(f"处理 {index_code} VS {index_market_standard}")
        # 提醒：get_index_monthly_return取的时间区间为[begin_date, end_date],
        #      get_index_return取得时间区间为[begin_date[0:7], end_date[0:7]]
        # monthly_return = get_index_monthly_return(begin_date, end_date, index_code)
        monthly_return = get_index_return(
            index_code=index_code,
            start_date=begin_date,
            end_date=end_date,
        )
        monthly_return['end_date'] = monthly_return['end_date'].apply(lambda x: x[0:7])
        monthly_return = monthly_return.rename(columns={'end_date': 'tj', 'daily_return': 'i_price'})
        tmp = monthly_return[index_code].std()
        sigma = tmp * math.sqrt(12)
        LOG.info(f'Code: {index_code}，要年化的sigma，从 {tmp} 到 {sigma}')

        beta, market_mean = get_beta(index_market_standard, monthly_return)

        market_risk_premium = market_mean - m_risk_free
        tmp = m_risk_free + beta * market_risk_premium
        e_r = (tmp + 1)**12 - 1
        LOG.info(f'Code: {index_code}，要年化的预期收益率，从 {tmp} 到 {e_r}')
        this_dict = {
            'begin_date': begin_date,
            'end_date': end_date,
            'expected_return': e_r,
            'expected_std': sigma,
            'beta': beta,
            'market_risk_premium': market_risk_premium
        }
        LOG.info(this_dict)
        ans_dict.update({index_code: this_dict})

    LOG.info(ans_dict)
    LOG.info('Done')
    return ans_dict
